{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SwinTransformer\n",
    "https://github.com/microsoft/Swin-Transformer/blob/main/models/swin_transformer.py"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "import itertools\n",
    "from typing import Optional, Sequence, Tuple, Type, Union\n",
    "\n",
    "import numpy as np\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import torch.nn.functional as F\n",
    "from torch.nn import LayerNorm\n",
    "\n",
    "from monai.networks.blocks import MLPBlock as Mlp\n",
    "from monai.networks.blocks import PatchEmbed\n",
    "from monai.networks.layers import DropPath, trunc_normal_\n",
    "from monai.utils import optional_import\n",
    "\n",
    "rearrange, _ = optional_import(\"einops\", name=\"rearrange\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [],
   "source": [
    "def window_partition(x, window_size):\n",
    "    \"\"\"window partition operation based on: \"Liu et al.,\n",
    "    Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n",
    "    <https://arxiv.org/abs/2103.14030>\"\n",
    "    https://github.com/microsoft/Swin-Transformer\n",
    "\n",
    "     Args:\n",
    "        x: input tensor.\n",
    "        window_size: local window size.\n",
    "    \"\"\"\n",
    "    b, h, w, c = x.shape\n",
    "    x = x.view(b, h // window_size[0], window_size[0], w // window_size[1], window_size[1], c)\n",
    "    windows = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(-1, window_size[0] * window_size[1], c)\n",
    "    return windows\n",
    "\n",
    "\n",
    "def window_reverse(windows, window_size, dims):\n",
    "    \"\"\"window reverse operation based on: \"Liu et al.,\n",
    "    Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n",
    "    <https://arxiv.org/abs/2103.14030>\"\n",
    "    https://github.com/microsoft/Swin-Transformer\n",
    "\n",
    "     Args:\n",
    "        windows: windows tensor.\n",
    "        window_size: local window size.\n",
    "        dims: dimension values.\n",
    "    \"\"\"\n",
    "    b, h, w = dims\n",
    "    x = windows.view(b, h // window_size[0], w // window_size[1], window_size[0], window_size[1], -1)\n",
    "    x = x.permute(0, 1, 3, 2, 4, 5).contiguous().view(b, h, w, -1)\n",
    "    return x\n",
    "\n",
    "\n",
    "def get_window_size(x_size, window_size, shift_size=None):\n",
    "    \"\"\"Computing window size based on: \"Liu et al.,\n",
    "    Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n",
    "    <https://arxiv.org/abs/2103.14030>\"\n",
    "    https://github.com/microsoft/Swin-Transformer\n",
    "\n",
    "     Args:\n",
    "        x_size: input size.\n",
    "        window_size: local window size.\n",
    "        shift_size: window shifting size.\n",
    "    \"\"\"\n",
    "\n",
    "    use_window_size = list(window_size)\n",
    "    if shift_size is not None:\n",
    "        use_shift_size = list(shift_size)\n",
    "    for i in range(len(x_size)):\n",
    "        if x_size[i] <= window_size[i]:\n",
    "            use_window_size[i] = x_size[i]\n",
    "            if shift_size is not None:\n",
    "                use_shift_size[i] = 0\n",
    "\n",
    "    if shift_size is None:\n",
    "        return tuple(use_window_size)\n",
    "    else:\n",
    "        return tuple(use_window_size), tuple(use_shift_size)\n",
    "\n",
    "\n",
    "class WindowAttention(nn.Module):\n",
    "    \"\"\"\n",
    "    Window based multi-head self attention module with relative position bias based on: \"Liu et al.,\n",
    "    Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n",
    "    <https://arxiv.org/abs/2103.14030>\"\n",
    "    https://github.com/microsoft/Swin-Transformer\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        dim: int,\n",
    "        num_heads: int,\n",
    "        window_size: Sequence[int],\n",
    "        qkv_bias: bool = False,\n",
    "        attn_drop: float = 0.0,\n",
    "        proj_drop: float = 0.0,\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            dim: number of feature channels.\n",
    "            num_heads: number of attention heads.\n",
    "            window_size: local window size.\n",
    "            qkv_bias: add a learnable bias to query, key, value.\n",
    "            attn_drop: attention dropout rate.\n",
    "            proj_drop: dropout rate of output.\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.window_size = window_size\n",
    "        self.num_heads = num_heads\n",
    "        head_dim = dim // num_heads\n",
    "        self.scale = head_dim**-0.5\n",
    "        mesh_args = torch.meshgrid.__kwdefaults__\n",
    "\n",
    "\n",
    "        self.relative_position_bias_table = nn.Parameter(\n",
    "            torch.zeros((2 * window_size[0] - 1) * (2 * window_size[1] - 1), num_heads)\n",
    "        )\n",
    "        coords_h = torch.arange(self.window_size[0])\n",
    "        coords_w = torch.arange(self.window_size[1])\n",
    "        if mesh_args is not None:\n",
    "            coords = torch.stack(torch.meshgrid(coords_h, coords_w, indexing=\"ij\"))\n",
    "        else:\n",
    "            coords = torch.stack(torch.meshgrid(coords_h, coords_w))\n",
    "        coords_flatten = torch.flatten(coords, 1)\n",
    "        relative_coords = coords_flatten[:, :, None] - coords_flatten[:, None, :]\n",
    "        relative_coords = relative_coords.permute(1, 2, 0).contiguous()\n",
    "        relative_coords[:, :, 0] += self.window_size[0] - 1\n",
    "        relative_coords[:, :, 1] += self.window_size[1] - 1\n",
    "        relative_coords[:, :, 0] *= 2 * self.window_size[1] - 1\n",
    "\n",
    "\n",
    "        relative_position_index = relative_coords.sum(-1)\n",
    "        self.register_buffer(\"relative_position_index\", relative_position_index)\n",
    "        self.qkv = nn.Linear(dim, dim * 3, bias=qkv_bias)\n",
    "        self.attn_drop = nn.Dropout(attn_drop)\n",
    "        self.proj = nn.Linear(dim, dim)\n",
    "        self.proj_drop = nn.Dropout(proj_drop)\n",
    "        trunc_normal_(self.relative_position_bias_table, std=0.02)\n",
    "        self.softmax = nn.Softmax(dim=-1)\n",
    "\n",
    "    def forward(self, x, mask):\n",
    "        b, n, c = x.shape\n",
    "        qkv = self.qkv(x).reshape(b, n, 3, self.num_heads, c // self.num_heads).permute(2, 0, 3, 1, 4)\n",
    "        q, k, v = qkv[0], qkv[1], qkv[2]\n",
    "        q = q * self.scale\n",
    "        attn = q @ k.transpose(-2, -1)\n",
    "        relative_position_bias = self.relative_position_bias_table[\n",
    "            self.relative_position_index.clone()[:n, :n].reshape(-1)\n",
    "        ].reshape(n, n, -1)\n",
    "        relative_position_bias = relative_position_bias.permute(2, 0, 1).contiguous()\n",
    "        attn = attn + relative_position_bias.unsqueeze(0)\n",
    "        if mask is not None:\n",
    "            nw = mask.shape[0]\n",
    "            attn = attn.view(b // nw, nw, self.num_heads, n, n) + mask.unsqueeze(1).unsqueeze(0)\n",
    "            attn = attn.view(-1, self.num_heads, n, n)\n",
    "            attn = self.softmax(attn)\n",
    "        else:\n",
    "            attn = self.softmax(attn)\n",
    "\n",
    "        attn = self.attn_drop(attn).to(v.dtype)\n",
    "        x = (attn @ v).transpose(1, 2).reshape(b, n, c)\n",
    "        x = self.proj(x)\n",
    "        x = self.proj_drop(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class SwinTransformerBlock(nn.Module):\n",
    "    \"\"\"\n",
    "    Swin Transformer block based on: \"Liu et al.,\n",
    "    Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n",
    "    <https://arxiv.org/abs/2103.14030>\"\n",
    "    https://github.com/microsoft/Swin-Transformer\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        dim: int,\n",
    "        num_heads: int,\n",
    "        window_size: Sequence[int],\n",
    "        shift_size: Sequence[int],\n",
    "        mlp_ratio: float = 4.0,\n",
    "        qkv_bias: bool = True,\n",
    "        drop: float = 0.0,\n",
    "        attn_drop: float = 0.0,\n",
    "        drop_path: float = 0.0,\n",
    "        act_layer: str = \"GELU\",\n",
    "        norm_layer: Type[LayerNorm] = nn.LayerNorm\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            dim: number of feature channels.\n",
    "            num_heads: number of attention heads.\n",
    "            window_size: local window size.\n",
    "            shift_size: window shift size.\n",
    "            mlp_ratio: ratio of mlp hidden dim to embedding dim.\n",
    "            qkv_bias: add a learnable bias to query, key, value.\n",
    "            drop: dropout rate.\n",
    "            attn_drop: attention dropout rate.\n",
    "            drop_path: stochastic depth rate.\n",
    "            act_layer: activation layer.\n",
    "            norm_layer: normalization layer.\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "        self.num_heads = num_heads\n",
    "        self.window_size = window_size\n",
    "        self.shift_size = shift_size\n",
    "        self.mlp_ratio = mlp_ratio\n",
    "        self.norm1 = norm_layer(dim)\n",
    "        self.attn = WindowAttention(\n",
    "            dim,\n",
    "            window_size=self.window_size,\n",
    "            num_heads=num_heads,\n",
    "            qkv_bias=qkv_bias,\n",
    "            attn_drop=attn_drop,\n",
    "            proj_drop=drop,\n",
    "        )\n",
    "\n",
    "        self.drop_path = DropPath(drop_path) if drop_path > 0.0 else nn.Identity()\n",
    "        self.norm2 = norm_layer(dim)\n",
    "        mlp_hidden_dim = int(dim * mlp_ratio)\n",
    "        self.mlp = Mlp(hidden_size=dim, mlp_dim=mlp_hidden_dim, act=act_layer, dropout_rate=drop, dropout_mode=\"swin\")\n",
    "\n",
    "    def forward_part1(self, x, mask_matrix):\n",
    "        x = self.norm1(x)\n",
    "\n",
    "        b, h, w, c = x.shape\n",
    "        window_size, shift_size = get_window_size((h, w), self.window_size, self.shift_size)\n",
    "        pad_l = pad_t = 0\n",
    "        pad_b = (window_size[0] - h % window_size[0]) % window_size[0]\n",
    "        pad_r = (window_size[1] - w % window_size[1]) % window_size[1]\n",
    "        x = F.pad(x, (0, 0, pad_l, pad_r, pad_t, pad_b))\n",
    "        _, hp, wp, _ = x.shape\n",
    "        dims = [b, hp, wp]\n",
    "\n",
    "        if any(i > 0 for i in shift_size):\n",
    "            shifted_x = torch.roll(x, shifts=(-shift_size[0], -shift_size[1]), dims=(1, 2))\n",
    "            attn_mask = mask_matrix\n",
    "        else:\n",
    "            shifted_x = x\n",
    "            attn_mask = None\n",
    "            \n",
    "        x_windows = window_partition(shifted_x, window_size)\n",
    "        attn_windows = self.attn(x_windows, mask=attn_mask)\n",
    "        attn_windows = attn_windows.view(-1, *(window_size + (c,)))\n",
    "        shifted_x = window_reverse(attn_windows, window_size, dims)\n",
    "        \n",
    "        if any(i > 0 for i in shift_size):\n",
    "            x = torch.roll(shifted_x, shifts=(shift_size[0], shift_size[1]), dims=(1, 2))\n",
    "        else:\n",
    "            x = shifted_x\n",
    "\n",
    "        if pad_r > 0 or pad_b > 0:\n",
    "            x = x[:, :h, :w, :].contiguous()\n",
    "            \n",
    "        return x\n",
    "\n",
    "    def forward_part2(self, x):\n",
    "        return self.drop_path(self.mlp(self.norm2(x)))\n",
    "\n",
    "    def forward(self, x, mask_matrix):\n",
    "        shortcut = x\n",
    "        x = self.forward_part1(x, mask_matrix)\n",
    "        x = shortcut + self.drop_path(x)\n",
    "        x = x + self.forward_part2(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class PatchMergingV2(nn.Module):\n",
    "    \"\"\"\n",
    "    Patch merging layer based on: \"Liu et al.,\n",
    "    Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n",
    "    <https://arxiv.org/abs/2103.14030>\"\n",
    "    https://github.com/microsoft/Swin-Transformer\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(self, dim: int, norm_layer: Type[LayerNorm] = nn.LayerNorm, spatial_dims: int = 3) -> None:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            dim: number of feature channels.\n",
    "            norm_layer: normalization layer.\n",
    "            spatial_dims: number of spatial dims.\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        self.dim = dim\n",
    "\n",
    "        self.reduction = nn.Linear(4 * dim, 2 * dim, bias=False)\n",
    "        self.norm = norm_layer(4 * dim)\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_shape = x.size()\n",
    "\n",
    "        b, h, w, c = x_shape\n",
    "        pad_input = (h % 2 == 1) or (w % 2 == 1)\n",
    "        if pad_input:\n",
    "            x = F.pad(x, (0, 0, 0, w % 2, 0, h % 2))\n",
    "        x = torch.cat([x[:, j::2, i::2, :] for i, j in itertools.product(range(2), range(2))], -1)\n",
    "\n",
    "        x = self.norm(x)\n",
    "        x = self.reduction(x)\n",
    "        return x\n",
    "\n",
    "\n",
    "class PatchMerging(PatchMergingV2):\n",
    "    \"\"\"The `PatchMerging` module previously defined in v0.9.0.\"\"\"\n",
    "    def forward(self, x):\n",
    "        return super().forward(x)\n",
    "\n",
    "\n",
    "MERGING_MODE = {\"merging\": PatchMerging, \"mergingv2\": PatchMergingV2}\n",
    "\n",
    "\n",
    "def compute_mask(dims, window_size, shift_size, device):\n",
    "    \"\"\"Computing region masks based on: \"Liu et al.,\n",
    "    Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n",
    "    <https://arxiv.org/abs/2103.14030>\"\n",
    "    https://github.com/microsoft/Swin-Transformer\n",
    "\n",
    "     Args:\n",
    "        dims: dimension values.\n",
    "        window_size: local window size.\n",
    "        shift_size: shift size.\n",
    "        device: device.\n",
    "    \"\"\"\n",
    "\n",
    "    cnt = 0\n",
    "\n",
    "    h, w = dims\n",
    "    img_mask = torch.zeros((1, h, w, 1), device=device)\n",
    "    for h in slice(-window_size[0]), slice(-window_size[0], -shift_size[0]), slice(-shift_size[0], None):\n",
    "        for w in slice(-window_size[1]), slice(-window_size[1], -shift_size[1]), slice(-shift_size[1], None):\n",
    "            img_mask[:, h, w, :] = cnt\n",
    "            cnt += 1\n",
    "\n",
    "    mask_windows = window_partition(img_mask, window_size)\n",
    "    mask_windows = mask_windows.squeeze(-1)\n",
    "    attn_mask = mask_windows.unsqueeze(1) - mask_windows.unsqueeze(2)\n",
    "    attn_mask = attn_mask.masked_fill(attn_mask != 0, float(-100.0)).masked_fill(attn_mask == 0, float(0.0))\n",
    "\n",
    "    return attn_mask\n",
    "\n",
    "\n",
    "class BasicLayer(nn.Module):\n",
    "    \"\"\"\n",
    "    Basic Swin Transformer layer in one stage based on: \"Liu et al.,\n",
    "    Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n",
    "    <https://arxiv.org/abs/2103.14030>\"\n",
    "    https://github.com/microsoft/Swin-Transformer\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        dim: int,\n",
    "        depth: int,\n",
    "        num_heads: int,\n",
    "        window_size: Sequence[int],\n",
    "        drop_path: list,\n",
    "        mlp_ratio: float = 4.0,\n",
    "        qkv_bias: bool = False,\n",
    "        drop: float = 0.0,\n",
    "        attn_drop: float = 0.0,\n",
    "        norm_layer: Type[LayerNorm] = nn.LayerNorm,\n",
    "        downsample: Optional[nn.Module] = None,\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            dim: number of feature channels.\n",
    "            depth: number of layers in each stage.\n",
    "            num_heads: number of attention heads.\n",
    "            window_size: local window size.\n",
    "            drop_path: stochastic depth rate.\n",
    "            mlp_ratio: ratio of mlp hidden dim to embedding dim.\n",
    "            qkv_bias: add a learnable bias to query, key, value.\n",
    "            drop: dropout rate.\n",
    "            attn_drop: attention dropout rate.\n",
    "            norm_layer: normalization layer.\n",
    "            downsample: an optional downsampling layer at the end of the layer.\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        self.window_size = window_size\n",
    "        self.shift_size = tuple(i // 2 for i in window_size)\n",
    "        self.no_shift = tuple(0 for i in window_size)\n",
    "        self.depth = depth\n",
    "        self.blocks = nn.ModuleList(\n",
    "            [\n",
    "                SwinTransformerBlock(\n",
    "                    dim=dim,\n",
    "                    num_heads=num_heads,\n",
    "                    window_size=self.window_size,\n",
    "                    shift_size=self.no_shift if (i % 2 == 0) else self.shift_size,\n",
    "                    mlp_ratio=mlp_ratio,\n",
    "                    qkv_bias=qkv_bias,\n",
    "                    drop=drop,\n",
    "                    attn_drop=attn_drop,\n",
    "                    drop_path=drop_path[i] if isinstance(drop_path, list) else drop_path,\n",
    "                    norm_layer=norm_layer,\n",
    "                )\n",
    "                for i in range(depth)\n",
    "            ]\n",
    "        )\n",
    "        self.downsample = downsample\n",
    "        if callable(self.downsample):\n",
    "            self.downsample = downsample(dim=dim, norm_layer=norm_layer, spatial_dims=len(self.window_size))\n",
    "\n",
    "    def forward(self, x):\n",
    "        x_shape = x.size()\n",
    "\n",
    "        b, c, h, w = x_shape\n",
    "        window_size, shift_size = get_window_size((h, w), self.window_size, self.shift_size)\n",
    "        x = rearrange(x, \"b c h w -> b h w c\")\n",
    "        hp = int(np.ceil(h / window_size[0])) * window_size[0]\n",
    "        wp = int(np.ceil(w / window_size[1])) * window_size[1]\n",
    "        attn_mask = compute_mask([hp, wp], window_size, shift_size, x.device)\n",
    "        \n",
    "        for blk in self.blocks:\n",
    "            x = blk(x, attn_mask)\n",
    "        x = x.view(b, h, w, -1)\n",
    "        \n",
    "        if self.downsample is not None:\n",
    "            x = self.downsample(x)\n",
    "        x = rearrange(x, \"b h w c -> b c h w\")\n",
    "        return x\n",
    "\n",
    "\n",
    "class SwinTransformer(nn.Module):\n",
    "    \"\"\"\n",
    "    Swin Transformer based on: \"Liu et al.,\n",
    "    Swin Transformer: Hierarchical Vision Transformer using Shifted Windows\n",
    "    <https://arxiv.org/abs/2103.14030>\"\n",
    "    https://github.com/microsoft/Swin-Transformer\n",
    "    \"\"\"\n",
    "\n",
    "    def __init__(\n",
    "        self,\n",
    "        in_chans: int,\n",
    "        embed_dim: int,\n",
    "        window_size: Sequence[int],\n",
    "        patch_size: Sequence[int],\n",
    "        depths: Sequence[int],\n",
    "        num_heads: Sequence[int],\n",
    "        mlp_ratio: float = 4.0,\n",
    "        qkv_bias: bool = True,\n",
    "        drop_rate: float = 0.0,\n",
    "        attn_drop_rate: float = 0.0,\n",
    "        drop_path_rate: float = 0.0,\n",
    "        norm_layer: Type[LayerNorm] = nn.LayerNorm,\n",
    "        patch_norm: bool = False,\n",
    "        spatial_dims: int = 3,\n",
    "        downsample=\"merging\",\n",
    "    ) -> None:\n",
    "        \"\"\"\n",
    "        Args:\n",
    "            in_chans: dimension of input channels.\n",
    "            embed_dim: number of linear projection output channels.\n",
    "            window_size: local window size.\n",
    "            patch_size: patch size.\n",
    "            depths: number of layers in each stage.\n",
    "            num_heads: number of attention heads.\n",
    "            mlp_ratio: ratio of mlp hidden dim to embedding dim.\n",
    "            qkv_bias: add a learnable bias to query, key, value.\n",
    "            drop_rate: dropout rate.\n",
    "            attn_drop_rate: attention dropout rate.\n",
    "            drop_path_rate: stochastic depth rate.\n",
    "            norm_layer: normalization layer.\n",
    "            patch_norm: add normalization after patch embedding.\n",
    "            spatial_dims: spatial dimension.\n",
    "            downsample: module used for downsampling, available options are `\"mergingv2\"`, `\"merging\"` and a\n",
    "                user-specified `nn.Module` following the API defined in :py:class:`monai.networks.nets.PatchMerging`.\n",
    "                The default is currently `\"merging\"` (the original version defined in v0.9.0).\n",
    "            phase: pretrain or finetune\n",
    "        \"\"\"\n",
    "\n",
    "        super().__init__()\n",
    "        self.num_layers = len(depths)\n",
    "        self.embed_dim = embed_dim\n",
    "        self.patch_norm = patch_norm\n",
    "        self.window_size = window_size\n",
    "        self.patch_size = patch_size\n",
    "\n",
    "        self.patch_embed = PatchEmbed(\n",
    "            patch_size=self.patch_size,\n",
    "            in_chans=in_chans,\n",
    "            embed_dim=embed_dim,\n",
    "            norm_layer=norm_layer if self.patch_norm else None,  # type: ignore\n",
    "            spatial_dims=spatial_dims,\n",
    "        )\n",
    "\n",
    "        self.pos_drop = nn.Dropout(p=drop_rate)\n",
    "        dpr = [x.item() for x in torch.linspace(0, drop_path_rate, sum(depths))]\n",
    "        self.layers1 = nn.ModuleList()\n",
    "        self.layers2 = nn.ModuleList()\n",
    "        self.layers3 = nn.ModuleList()\n",
    "        self.layers4 = nn.ModuleList()\n",
    "        down_sample_mod = look_up_option(downsample, MERGING_MODE) if isinstance(downsample, str) else downsample\n",
    "        for i_layer in range(self.num_layers):\n",
    "            layer = BasicLayer(\n",
    "                dim=int(embed_dim * 2**i_layer),\n",
    "                depth=depths[i_layer],\n",
    "                num_heads=num_heads[i_layer],\n",
    "                window_size=self.window_size,\n",
    "                drop_path=dpr[sum(depths[:i_layer]) : sum(depths[: i_layer + 1])],\n",
    "                mlp_ratio=mlp_ratio,\n",
    "                qkv_bias=qkv_bias,\n",
    "                drop=drop_rate,\n",
    "                attn_drop=attn_drop_rate,\n",
    "                norm_layer=norm_layer,\n",
    "                downsample=down_sample_mod,\n",
    "            )\n",
    "            if i_layer == 0:\n",
    "                self.layers1.append(layer)\n",
    "            elif i_layer == 1:\n",
    "                self.layers2.append(layer)\n",
    "            elif i_layer == 2:\n",
    "                self.layers3.append(layer)\n",
    "            elif i_layer == 3:\n",
    "                self.layers4.append(layer)\n",
    "        self.num_features = int(embed_dim * 2 ** (self.num_layers - 1))\n",
    "\n",
    "    def proj_out(self, x, normalize=False):\n",
    "        if normalize:\n",
    "            x_shape = x.size()\n",
    "            \n",
    "            n, ch, h, w = x_shape\n",
    "            x = rearrange(x, \"n c h w -> n h w c\")\n",
    "            x = F.layer_norm(x, [ch])\n",
    "            x = rearrange(x, \"n h w c -> n c h w\")\n",
    "        return x\n",
    "\n",
    "    def forward(self, x, normalize=True):\n",
    "        x0 = self.patch_embed(x)\n",
    "        x0 = self.pos_drop(x0)\n",
    "\n",
    "        x1 = self.layers1[0](x0.contiguous())\n",
    "        x2 = self.layers2[0](x1.contiguous())\n",
    "        x3 = self.layers3[0](x2.contiguous())\n",
    "        x4 = self.layers4[0](x3.contiguous())\n",
    "        x4_out = self.proj_out(x4, normalize)\n",
    "        \n",
    "        # todo add linear projector\n",
    "        # x = self.avgpool(x.transpose(1, 2))  # B C 1\n",
    "        # x = torch.flatten(x, 1)\n",
    "\n",
    "        return x4_out"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "torch.Size([2, 384, 8, 8])\n"
     ]
    }
   ],
   "source": [
    "swinViT = SwinTransformer(\n",
    "    in_chans=4,\n",
    "    embed_dim=24,\n",
    "    window_size=(7, 7),\n",
    "    patch_size=(2, 2),\n",
    "    depths=(2, 2, 2, 2),\n",
    "    num_heads=(3, 6, 12, 24),\n",
    "    mlp_ratio=4.0,\n",
    "    qkv_bias=True,\n",
    "    drop_rate=0.0,\n",
    "    attn_drop_rate=0.0,\n",
    "    drop_path_rate=0.0,\n",
    "    norm_layer=nn.LayerNorm,\n",
    "    spatial_dims=2,\n",
    "    downsample=look_up_option(\"merging\", MERGING_MODE)\n",
    ")\n",
    "\n",
    "inps_2D = torch.randn((2, 4, 256, 256))\n",
    "\n",
    "out_2D = swinViT(inps_2D)\n",
    "\n",
    "print(out_2D.size())"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "swinuneter",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
